{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/DeepLabCut/DeepLabCut/blob/master/examples/COLAB/COLAB_maDLC_TrainNetwork_VideoAnalysis.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RK255E7YoEIt"
   },
   "source": [
    "# DeepLabCut for your multi-animal projects!\n",
    "\n",
    "Some useful links:\n",
    "\n",
    "- [DeepLabCut's GitHub: github.com/DeepLabCut/DeepLabCut](https://github.com/DeepLabCut/DeepLabCut)\n",
    "- [DeepLabCut's Documentation: User Guide for Multi-Animal projects](https://deeplabcut.github.io/DeepLabCut/docs/maDLC_UserGuide.html)\n",
    "\n",
    "\n",
    "![alt text](https://images.squarespace-cdn.com/content/v1/57f6d51c9f74566f55ecf271/1628180434489-T0RIWEJJU0FJVOT6FNVD/maDLC.png?format=800w)\n",
    "\n",
    "This notebook illustrates how to, for multi-animal projects, use the cloud-based GPU to:\n",
    "- create a multi-animal training set\n",
    "- train a network\n",
    "- evaluate a network\n",
    "- analyze novel videos\n",
    "- assemble animals and tracklets\n",
    "- create quality check plots!\n",
    "\n",
    "### This notebook assumes you already have a DLC project folder with labeled data and you uploaded it to your own Google Drive.\n",
    "\n",
    "This notebook demonstrates the necessary steps to use DeepLabCut for your own project.\n",
    "\n",
    "This shows the most simple code to do so, but many of the functions have additional features, so please check out the docs on GitHub. We also recommend checking out our preprint, which covers the science of maDLC\n",
    "\n",
    "**Lauer et al 2021:** https://www.biorxiv.org/content/10.1101/2021.04.30.442096v1\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "txoddlM8hLKm"
   },
   "source": [
    "## First, go to \"Runtime\" ->\"change runtime type\"->select \"Python3\", and then select \"GPU\"\n",
    "\n",
    "As the COLAB environments were updated to CUDA 12.X and Python 3.11, we need to install DeepLabCut and TensorFlow in a distinct way to get TensorFlow to connect to the GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this will take a couple of minutes to install all the dependencies!\n",
    "!pip install --pre deeplabcut"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**(Be sure to click \"RESTART RUNTIME\" if it is displayed above before moving on !)** You will see this button at the output of the cells above ^."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "oTwAcbq2-FZz",
    "outputId": "9cfd8dcf-a0a8-4801-ed1d-fbcd5ec056af"
   },
   "outputs": [],
   "source": [
    "import deeplabcut"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cQ-nlTkri4HZ"
   },
   "source": [
    "## Link your Google Drive (with your labeled data):\n",
    "\n",
    "- This code assumes you locally installed DeepLabCut, created a project, extracted and labeled frames. Be sure to \"check Labels\" to confirm you are happy with your data. As, these frames are the only thing that is used to train your network. 💪 You can find all the docs to do this here: [deeplabcut.github.io/DeepLabCut](https://deeplabcut.github.io/DeepLabCut/README.html)\n",
    "- Next, place your DLC project folder into you Google Drive- i.e., copy the folder named \"Project-YourName-TheDate\" into Google Drive.\n",
    "- Then, click run on the cell below to link this notebook to your Google Drive:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KS4Q4UkR9rgG"
   },
   "outputs": [],
   "source": [
    "# Now, let's link to your GoogleDrive. Run this cell and follow the authorization instructions:\n",
    "# (We recommend putting a copy of the github repo in your google drive if you are using the demo \"examples\")\n",
    "\n",
    "from google.colab import drive\n",
    "\n",
    "drive.mount(\"/content/drive\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Frnj1RVDyEqs"
   },
   "source": [
    "## Next, edit the few items below, and click run:\n",
    "\n",
    "YOU WILL NEED TO EDIT THE PROJECT PATH **in the `config.yaml` file** TO BE SET TO YOUR GOOGLE DRIVE LINK! Typically, this will be: `/content/drive/My Drive/yourProjectFolderName`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vhENAlQnFENJ"
   },
   "outputs": [],
   "source": [
    "# PLEASE EDIT THIS:\n",
    "project_folder_name = \"MontBlanc-Daniel-2019-12-16\"\n",
    "video_type = \"mp4\" #, mp4, MOV, or avi, whatever you uploaded!\n",
    "\n",
    "# No need to edit this, we are going to assume you put videos you want to analyze\n",
    "# in the \"videos\" folder, but if this is NOT true, edit below:\n",
    "videofile_path = [f\"/content/drive/My Drive/{project_folder_name}/videos/\"]\n",
    "print(videofile_path)\n",
    "\n",
    "# The prediction files and labeled videos will be saved in this `labeled-videos` folder\n",
    "# in your project folder; if you want them elsewhere, you can edit this;\n",
    "# if you want the output files in the same folder as the videos, set this to an empty string.\n",
    "destfolder = f\"/content/drive/My Drive/{project_folder_name}/labeled-videos\"\n",
    "\n",
    "#No need to edit this, as you set it when you passed the ProjectFolderName (above):\n",
    "path_config_file = f\"/content/drive/My Drive/{project_folder_name}/config.yaml\"\n",
    "print(path_config_file)\n",
    "\n",
    "# This creates a path variable that links to your Google Drive project"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xNi9s1dboEJN"
   },
   "source": [
    "## Create a multi-animal training dataset:\n",
    "\n",
    "- more info can be [found in the docs](https://deeplabcut.github.io/DeepLabCut/docs/maDLC_UserGuide.html#create-training-dataset)\n",
    "- please check the text below, edit if needed, and then click run (this can take some time):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-AAYSoW313me"
   },
   "outputs": [],
   "source": [
    "# OPTIONAL LEARNING: did you know you can check what each function does by running with a ?\n",
    "deeplabcut.create_multianimaltraining_dataset?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eMeUwgxPoEJP",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# ATTENTION:\n",
    "# Which shuffle do you want to create and train?\n",
    "shuffle = 1 # Edit if needed; 1 is the default.\n",
    "\n",
    "deeplabcut.create_multianimaltraining_dataset(\n",
    "    path_config_file,\n",
    "    Shuffles=[shuffle],\n",
    "    net_type=\"dlcrnet_ms5\",\n",
    "    engine=deeplabcut.Engine.PYTORCH,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "c4FczXGDoEJU"
   },
   "source": [
    "## Start training:\n",
    "This function trains the network for a specific shuffle of the training dataset. More info can be found [in the docs](https://deeplabcut.github.io/DeepLabCut/docs/maDLC_UserGuide.html#train-the-network)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_pOvDq_2oEJW"
   },
   "outputs": [],
   "source": [
    "# Let's also change the display and save_epochs just in case Colab takes away\n",
    "# the GPU... If that happens, you can reload from a saved point using the\n",
    "# `snapshot_path` argument to `deeplabcut.train_network`:\n",
    "#   deeplabcut.train_network(..., snapshot_path=\"/content/.../snapshot-050.pt\")\n",
    "\n",
    "# Typically, you want to train to ~200 epochs. We set the batch size to 8 to\n",
    "# utilize the GPU's capabilities.\n",
    "\n",
    "# More info and there are more things you can set:\n",
    "#   https://deeplabcut.github.io/DeepLabCut/docs/standardDeepLabCut_UserGuide.html#g-train-the-network\n",
    "\n",
    "deeplabcut.train_network(\n",
    "    path_config_file,\n",
    "    shuffle=shuffle,\n",
    "    save_epochs=5,\n",
    "    epochs=200,\n",
    "    batch_size=8,\n",
    ")\n",
    "\n",
    "# This will run until you stop it (CTRL+C), or hit \"STOP\" icon, or when it hits the end."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RiDwIVf5-3H_"
   },
   "source": [
    "Note, that **when you hit \"STOP\" you will get a `KeyboardInterrupt` \"error\"! No worries! :)**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xZygsb2DoEJc"
   },
   "source": [
    "## Start evaluating: \n",
    "\n",
    "- First, we evaluate the pose estimation performance.\n",
    "- This function evaluates a trained model for a specific shuffle/shuffles at a particular state or all the states on the data set (images) and stores the results as .5 and .csv file in a subdirectory under **evaluation-results-pytorch**\n",
    "- If the scoremaps do not look accurate, don't proceed to tracklet assembly; please consider (1) adding more data, (2) adding more bodyparts!\n",
    "- More info can be [found in the docs](https://deeplabcut.github.io/DeepLabCut/docs/maDLC_UserGuide.html#evaluate-the-trained-network)\n",
    "\n",
    "Here is an example of what you'd aim to see before proceeding:\n",
    "\n",
    "![alt text](https://images.squarespace-cdn.com/content/v1/57f6d51c9f74566f55ecf271/1590535809087-X655WY9W1MW1MY1I7DHE/ke17ZwdGBToddI8pDm48kBoswZhKnUtAF7-bTXgw67EUqsxRUqqbr1mOJYKfIPR7LoDQ9mXPOjoJoqy81S2I8N_N4V1vUb5AoIIIbLZhVYxCRW4BPu10St3TBAUQYVKc5tTP1cnANTUwNNPnYFjIp6XbP9N1GxIgAkxvBVqt0UvLpPHYwvNQTwHg8f_Zu8ZF/evaluation.png?format=1000w)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nv4zlbrnoEJg"
   },
   "outputs": [],
   "source": [
    "# Let's evaluate first:\n",
    "deeplabcut.evaluate_network(path_config_file, Shuffles=[shuffle], plotting=True)\n",
    "\n",
    "# plot a few scoremaps:\n",
    "deeplabcut.extract_save_all_maps(path_config_file, shuffle=shuffle, Indices=[0, 1, 2, 3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fYlGbloolDU2"
   },
   "source": [
    "IF these images, numbers, and maps do not look good, do not proceed. You should increase the diversity and number of frames you label, and re-create a training dataset and re-train! "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OVFLSKKfoEJk"
   },
   "source": [
    "## Start Analyzing videos: \n",
    "This function analyzes the new video. The user can choose the best model from the evaluation results and specify the correct snapshot index for the variable **snapshotindex** in the **config.yaml** file. Otherwise, by default the most recent snapshot is used to analyse the video.\n",
    "\n",
    "The results are stored in a pickle file in the same directory where the video resides. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Y_LZiS_0oEJl"
   },
   "outputs": [],
   "source": [
    "print(\"Start Analyzing my video(s)!\")\n",
    "#EDIT OPTION: which video(s) do you want to analyze? You can pass a path or a folder:\n",
    "# currently, if you run \"as is\" it assumes you have a video in the DLC project video folder!\n",
    "\n",
    "deeplabcut.analyze_videos(\n",
    "    path_config_file,\n",
    "    videofile_path,\n",
    "    shuffle=shuffle,\n",
    "    videotype=video_type,\n",
    "    auto_track=False,\n",
    "    destfolder=destfolder,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "91xBLOcBzGxo"
   },
   "source": [
    "Optional: Now you have the option to check the raw detections before animals are tracked. To do so, pass a video path:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "65mWwX5bTc5C"
   },
   "outputs": [],
   "source": [
    "##### PROTIP: #####\n",
    "## look at the output video; if the pose estimation (i.e. key points)\n",
    "## don't look good, don't proceed with tracking - add more data to your training set and re-train!\n",
    "\n",
    "# EDIT: let's check a specific video (PLEASE EDIT VIDEO PATH):\n",
    "specific_videofile = \"/content/drive/MyDrive/DeepLabCut_maDLC_DemoData/MontBlanc-Daniel-2019-12-16/videos/short.mov\"\n",
    "\n",
    "# Don't edit:\n",
    "deeplabcut.create_video_with_all_detections(\n",
    "    path_config_file, [specific_videofile], shuffle=shuffle, destfolder=destfolder,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3-OgTJ0Lz20e"
   },
   "source": [
    "If the resulting video (ends in full.mp4) is not good, we highly recommend adding more data and training again. See [here, in the docs](https://deeplabcut.github.io/DeepLabCut/docs/maDLC_UserGuide.html#decision-break-point)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PxRLS2_-r55K"
   },
   "source": [
    "## Next, we will assemble animals using our data-driven optimal graph method:\n",
    "\n",
    "During video analysis, animals are assembled using the optimal graph, which matches the \"data-driven\" method from our paper (Figure adapted from Lauer et al. 2021)\n",
    "\n",
    "![alt text](https://images.squarespace-cdn.com/content/v1/57f6d51c9f74566f55ecf271/1626266017809-XO6NX84QB4FBAZGOTCEY/fig3.jpg?format=400w)\n",
    "\n",
    "The optimal graph is computed when `evaluate_network` - so make sure you don't skip that step!\n",
    "\n",
    "**Note**: you can set the number of animals you expect to see, so check, edit, then click run:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zIvXM7TXIs-U"
   },
   "outputs": [],
   "source": [
    "#Check and edit:\n",
    "num_animals = 4 # How many animals do you expect to find?\n",
    "track_type= \"box\" # box, skeleton, ellipse\n",
    "#-- ellipse is recommended, unless you have a single-point MA project, then use BOX!\n",
    "\n",
    "# Optional:\n",
    "# imagine you tracked a point that is not useful for assembly,\n",
    "# like a tail tip that is far from the body, consider dropping it for this step (it's still used later)!\n",
    "# To drop it, uncomment the next line TWO lines and add your parts(s):\n",
    "\n",
    "# bodypart= 'Tail_end'\n",
    "# deeplabcut.convert_detections2tracklets(path_config_file, videofile_path, videotype=VideoType, shuffle=shuffle, overwrite=True, ignore_bodyparts=[bodypart])\n",
    "\n",
    "# OR don't drop, just click RUN:\n",
    "deeplabcut.convert_detections2tracklets(\n",
    "    path_config_file,\n",
    "    videofile_path,\n",
    "    videotype=video_type,\n",
    "    shuffle=shuffle,\n",
    "    track_method=track_type,\n",
    "    destfolder=destfolder,\n",
    "    overwrite=True,\n",
    ")\n",
    "\n",
    "deeplabcut.stitch_tracklets(\n",
    "    path_config_file,\n",
    "    videofile_path,\n",
    "    shuffle=shuffle,\n",
    "    track_method=track_type,\n",
    "    n_tracks=num_animals,\n",
    "    destfolder=destfolder,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TqbAnyfL0Q7h"
   },
   "source": [
    "Now let's filter the data to remove any small jitter:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "a6izVWX8sdzL"
   },
   "outputs": [],
   "source": [
    "deeplabcut.filterpredictions(\n",
    "    path_config_file,\n",
    "    videofile_path,\n",
    "    shuffle=shuffle,\n",
    "    videotype=video_type,\n",
    "    track_method=track_type,\n",
    "    destfolder=destfolder,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Zk4xGb8Ftf3B"
   },
   "source": [
    "## Create plots of your trajectories:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gX21zZbXoEKJ"
   },
   "outputs": [],
   "source": [
    "deeplabcut.plot_trajectories(\n",
    "    path_config_file,\n",
    "    videofile_path,\n",
    "    videotype=video_type,\n",
    "    shuffle=shuffle,\n",
    "    track_method=track_type,\n",
    "    destfolder=destfolder,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pqaCw15v8EmB"
   },
   "source": [
    "Now you can look at the plot-poses file and check the \"plot-likelihood.png\" might want to change the \"p-cutoff\" in the config.yaml file so that you have only high confidnece points plotted in the video. i.e. ~0.8 or 0.9. The current default is 0.4. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pCrUvQIvoEKD"
   },
   "source": [
    "## Create labeled video:\n",
    "This function is for visualization purpose and can be used to create a video in .mp4 format with labels predicted by the network. This video is saved in the same directory where the original video resides. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6aDF7Q7KoEKE"
   },
   "outputs": [],
   "source": [
    "deeplabcut.create_labeled_video(\n",
    "    path_config_file,\n",
    "    videofile_path,\n",
    "    shuffle=shuffle,\n",
    "    color_by=\"individual\",\n",
    "    videotype=video_type,\n",
    "    save_frames=False,\n",
    "    filtered=True,\n",
    "    track_method=track_type,\n",
    "    destfolder=destfolder,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "COLAB_maDLC_TrainNetwork_VideoAnalysis.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
